import numpy as np
import matplotlib.pyplot as plt
import torch
# python image library python图像处理库
from PIL import Image
import torchvision
from torchvision import transforms, datasets
from torch.utils.data import DataLoader

p = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])

train = datasets.MNIST('../data', transform=p, train=True, download=True)
test = datasets.MNIST('../data', transform=p, train=False, download=True)
# print(d[1888])
# print(train[666][0].min())
# imgx = transforms.ToPILImage()(train[666][0])
# plt.imshow(imgx)
# plt.show()

train_loader = DataLoader(train, batch_size=16, shuffle=True)
test_loader = DataLoader(test, batch_size=16, shuffle=False)

# 对一个batch中的图片可视化
# 一:挤压掉c
# 二：chw --> hwc
samples = iter(train_loader)
images, labels = next(samples)
# for i in range(16):
#     plt.subplot(4, 4, i+1)
#     plt.tight_layout()
#     # 关掉坐标
#     plt.axis('off')
#     # 加标题
#     plt.title(labels[i].item())
#     plt.imshow(images[i].permute(1, 2, 0), cmap=plt.cm.gray)
# plt.show()

# 三
# def show_img(img):
#     img=img.numpy()
#     img = (img+1)/2
#     # img = img/2 + 0.5
#     img_np = np.transpose(img, (1, 2, 0))
#     plt.imshow(img_np)
#     plt.show()
#

# # 四
# plt.imshow(transforms.ToPILImage()(torchvision.utils.make_grid(images[:12], nrow=4, padding=18)))
# plt.show()

# 五 爱因斯坦求和法交换数据
a=torch.size([3, 4, 5, 6])
torch.einsum('abcd->cadb', a)

# show_img(images)























# 改变维度
# b = torch.